"""
AutoGen 股票分析智能体协作系统 - 主程序
展示 AutoGen 的核心功能和最佳实践
"""
from autogen_flow import AutoGenStockFlow
from dotenv import load_dotenv
from colorama import init, Fore, Style
import json
import os
import sys
from typing import Dict, Any


def display_autogen_features():
    """展示 AutoGen 功能特性"""
    print(Fore.CYAN + Style.BRIGHT + "\n🔥 === AutoGen 功能展示 ===")
    features = [
        "🤖 多智能体协作 - 自动化的智能体对话和协作",
        "💬 群聊管理 - GroupChat 支持多智能体群聊",
        "🔄 对话模式 - 支持顺序、并行、群聊等多种模式",
        "🛠️ 函数调用 - 原生支持工具和函数调用",
        "📝 代码执行 - 可选的代码执行能力",
        "🧠 上下文管理 - 自动的对话历史和上下文管理",
        "⚡ 灵活配置 - 高度可配置的智能体行为",
        "🎯 任务导向 - 面向任务的智能体设计"
    ]
    
    for feature in features:
        print(Fore.YELLOW + f"  {feature}")
    
    print(Fore.CYAN + Style.BRIGHT + "\n本项目展示了 AutoGen 在股票分析场景的应用！")


def check_autogen_environment():
    """检查 AutoGen 环境配置"""
    print(Fore.CYAN + "\n🔍 检查 AutoGen 环境...")
    
    try:
        # 检查 AutoGen 相关包
        import autogen
        print(Fore.GREEN + "✅ AutoGen 已安装")
        
        # 检查 OpenAI 包
        import openai
        print(Fore.GREEN + "✅ OpenAI 包已安装")
        
        # 检查配置文件
        config_path = os.path.join(os.path.dirname(__file__), "config.json")
        if os.path.exists(config_path):
            print(Fore.GREEN + "✅ 配置文件存在")
        else:
            print(Fore.RED + "❌ 配置文件缺失")
            return False
        
        # 检查环境变量
        load_dotenv()
        
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
        
        # 检查 LLM 配置
        llm_config = config.get("llm_config", {})
        config_list = llm_config.get("config_list", [])
        if config_list:
            for cfg in config_list:
                if cfg.get("api_key") and not cfg["api_key"].startswith("your_"):
                    print(Fore.GREEN + f"✅ {cfg.get('model', 'Unknown')} API Key 已配置")
                else:
                    print(Fore.YELLOW + f"⚠️ {cfg.get('model', 'Unknown')} API Key 需要配置")
        
        # 检查 Tushare Token
        tushare_token = config.get("tushare_token")
        if tushare_token and not tushare_token.startswith("your_"):
            print(Fore.GREEN + "✅ Tushare Token 已配置")
        else:
            print(Fore.YELLOW + "⚠️ Tushare Token 需要配置")
        
        return True
        
    except ImportError as e:
        print(Fore.RED + f"❌ AutoGen 包缺失: {e}")
        print(Fore.YELLOW + "请运行: pip install pyautogen")
        return False
    except Exception as e:
        print(Fore.RED + f"❌ 环境检查失败: {e}")
        return False


def interactive_demo():
    """交互式演示"""
    print(Fore.CYAN + Style.BRIGHT + "\n🎮 === 交互式演示模式 ===")
    
    autogen_flow = None
    
    while True:
        print(Fore.CYAN + "\n请选择演示选项:")
        print("1. 🚀 顺序执行模式（推荐）")
        print("2. 💬 群聊协作模式")
        print("3. 👥 两智能体对话模式")
        print("4. 📊 查看智能体信息")
        print("5. 📈 分析自定义股票")
        print("6. 🎯 查看 AutoGen 特性")
        print("7. 📋 框架对比")
        print("8. 🏃 退出演示")
        
        choice = input(Fore.GREEN + "\n请输入选项 (1-8): ").strip()
        
        if choice == "1":
            run_sequential_mode()
        elif choice == "2":
            run_group_chat_mode()
        elif choice == "3":
            run_two_agent_mode()
        elif choice == "4":
            show_agent_info()
        elif choice == "5":
            analyze_custom_stock()
        elif choice == "6":
            display_autogen_features()
        elif choice == "7":
            compare_frameworks()
        elif choice == "8":
            print(Fore.CYAN + "👋 感谢使用 AutoGen 智能体协作系统！")
            break
        else:
            print(Fore.RED + "❌ 无效选项，请重新选择")


def run_sequential_mode():
    """运行顺序执行模式"""
    print(Fore.CYAN + "\n🚀 启动 AutoGen 顺序执行模式...")
    
    try:
        autogen_flow = AutoGenStockFlow()
        stock_code = "600519.SH"  # 贵州茅台
        print(Fore.GREEN + f"📈 分析股票: {stock_code} (贵州茅台)")
        autogen_flow.run_sequential(stock_code)
        
    except Exception as e:
        print(Fore.RED + f"❌ 顺序执行失败: {e}")


def run_group_chat_mode():
    """运行群聊协作模式"""
    print(Fore.CYAN + "\n💬 启动 AutoGen 群聊协作模式...")
    
    try:
        autogen_flow = AutoGenStockFlow()
        stock_code = "600519.SH"
        print(Fore.GREEN + f"📈 分析股票: {stock_code} (贵州茅台)")
        autogen_flow.run_group_chat(stock_code)
        
    except Exception as e:
        print(Fore.RED + f"❌ 群聊协作失败: {e}")


def run_two_agent_mode():
    """运行两智能体对话模式"""
    print(Fore.CYAN + "\n👥 启动两智能体对话模式...")
    
    try:
        autogen_flow = AutoGenStockFlow()
        stock_code = "600519.SH"
        print(Fore.GREEN + f"📈 分析股票: {stock_code} (贵州茅台)")
        autogen_flow.run_two_agent_chat(stock_code)
        
    except Exception as e:
        print(Fore.RED + f"❌ 两智能体对话失败: {e}")


def show_agent_info():
    """显示智能体信息"""
    print(Fore.CYAN + "\n📊 AutoGen 智能体信息...")
    
    try:
        autogen_flow = AutoGenStockFlow()
        autogen_flow.get_agent_info()
        
    except Exception as e:
        print(Fore.RED + f"❌ 获取智能体信息失败: {e}")


def analyze_custom_stock():
    """分析自定义股票"""
    print(Fore.CYAN + "\n📈 自定义股票分析...")
    
    stock_code = input(Fore.GREEN + "请输入股票代码 (例如: 000001.SZ): ").strip()
    
    if not stock_code:
        print(Fore.RED + "❌ 股票代码不能为空")
        return
    
    print(Fore.CYAN + "\n请选择执行模式:")
    print("1. 顺序执行")
    print("2. 群聊协作")
    print("3. 两智能体对话")
    
    mode = input(Fore.GREEN + "请输入选项 (1-3): ").strip()
    
    try:
        autogen_flow = AutoGenStockFlow()
        
        if mode == "1":
            autogen_flow.run_sequential(stock_code)
        elif mode == "2":
            autogen_flow.run_group_chat(stock_code)
        elif mode == "3":
            autogen_flow.run_two_agent_chat(stock_code)
        else:
            print(Fore.YELLOW + "使用默认顺序执行模式...")
            autogen_flow.run_sequential(stock_code)
            
    except Exception as e:
        print(Fore.RED + f"❌ 股票分析失败: {e}")


def compare_frameworks():
    """框架对比"""
    print(Fore.CYAN + Style.BRIGHT + "\n📊 === 三大框架对比 ===")
    
    comparison = """
    ┌─────────────┬────────────────┬────────────────┬────────────────┐
    │   特性      │   LangChain    │    CrewAI      │   AutoGen      │
    ├─────────────┼────────────────┼────────────────┼────────────────┤
    │ 设计理念    │ 链式组合       │ 任务驱动       │ 对话驱动       │
    │ 核心概念    │ Chain/Runnable │ Agent/Task     │ Agent/Chat     │
    │ 协作方式    │ 管道式         │ 任务依赖       │ 对话协作       │
    │ 工具支持    │ Tool接口       │ 内置工具框架   │ 函数调用       │
    │ 代码执行    │ 需要扩展       │ 不支持         │ 原生支持       │
    │ 群聊支持    │ 需自定义       │ 不支持         │ GroupChat      │
    │ 学习曲线    │ 陡峭           │ 平缓           │ 适中           │
    │ 灵活性      │ 高             │ 中             │ 高             │
    │ 生态系统    │ 最丰富         │ 发展中         │ 快速增长       │
    │ 适用场景    │ 通用AI应用     │ 多智能体任务   │ 对话式AI系统   │
    └─────────────┴────────────────┴────────────────┴────────────────┘
    """
    print(Fore.YELLOW + comparison)
    
    print(Fore.CYAN + "\n💡 选择建议：")
    print(Fore.WHITE + "• LangChain: 需要灵活性和丰富生态的复杂应用")
    print(Fore.WHITE + "• CrewAI: 快速构建多智能体任务系统")
    print(Fore.WHITE + "• AutoGen: 需要对话协作和代码执行的场景")
    
    print(Fore.CYAN + "\n📈 本项目对比：")
    print(Fore.WHITE + "• 代码量: LangChain > AutoGen > CrewAI")
    print(Fore.WHITE + "• 易理解: CrewAI > AutoGen > LangChain")
    print(Fore.WHITE + "• 功能性: AutoGen ≈ LangChain > CrewAI")


def main():
    """主函数"""
    init(autoreset=True)
    
    print(Fore.CYAN + Style.BRIGHT + "=" * 60)
    print(Fore.CYAN + Style.BRIGHT + "🤖 AutoGen 股票分析智能体协作系统")
    print(Fore.CYAN + Style.BRIGHT + "=" * 60)
    
    # 显示 AutoGen 功能特性
    display_autogen_features()
    
    # 检查环境
    if not check_autogen_environment():
        print(Fore.RED + "\n❌ 环境检查失败，请检查配置后重试")
        print(Fore.YELLOW + "\n请确保：")
        print(Fore.YELLOW + "1. 安装 AutoGen: pip install pyautogen")
        print(Fore.YELLOW + "2. 配置 API Key 在 config.json")
        print(Fore.YELLOW + "3. 配置 Tushare Token")
        return
    
    print(Fore.GREEN + "\n✅ 环境检查通过，系统准备就绪！")
    
    # 询问运行模式
    print(Fore.CYAN + "\n请选择运行模式:")
    print("1. 🚀 快速演示 (顺序执行)")
    print("2. 💬 群聊协作演示")
    print("3. 🎮 交互式演示")
    print("4. 📊 框架对比")
    
    mode = input(Fore.GREEN + "请输入选项 (1-4): ").strip()
    
    if mode == "1":
        print(Fore.CYAN + "\n🚀 启动快速演示模式...")
        run_sequential_mode()
    elif mode == "2":
        print(Fore.CYAN + "\n💬 启动群聊协作模式...")
        run_group_chat_mode()
    elif mode == "3":
        interactive_demo()
    elif mode == "4":
        compare_frameworks()
        # 显示对比后询问是否运行演示
        choice = input(Fore.GREEN + "\n是否运行 AutoGen 演示? (y/n): ").strip().lower()
        if choice == 'y':
            run_sequential_mode()
    else:
        print(Fore.YELLOW + "使用默认快速演示模式...")
        run_sequential_mode()
    
    print(Fore.CYAN + Style.BRIGHT + "\n🎉 感谢使用 AutoGen 智能体协作系统！")
    print(Fore.CYAN + "\n💡 提示：AutoGen 特别适合需要智能体对话和代码执行的场景")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print(Fore.YELLOW + "\n\n👋 用户中断，程序退出")
    except Exception as e:
        print(Fore.RED + f"\n❌ 程序异常退出: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
